#pragma once
#include <vector>
#include <fmt/format.h>
#include <fmt/ranges.h>
#include "./Entity3D.hpp"
#include "./Material.hpp"
#include "./Pipeline.hpp"

struct MeshGeometry {
    template <typename T> using array = std::vector<T>;

    array<glm::vec3> positions{};
    array<glm::vec3> normals{};
    array<glm::vec2> texCoords{};
    array<opengl::u32> indices{};
};

template <typename Material> class Mesh3D : public Entity3D {
public:
    Mesh3D(){};
    Mesh3D(const MeshGeometry &data, const Material &mat) : material(mat) { create(data); };
    inline void create(const MeshGeometry &data);

public:
    virtual void update();
    virtual void render();

public:
    Material material;
    opengl::buffer materialBuffer{sizeof(material)};
    void updateMaterial() { // Require to be called explicitly
        materialBuffer.set_data(&material, sizeof(material));
    }

public:
    Pipeline pipeline;
    opengl::program program{};

public:
    opengl::u32 numIndices{0};
    opengl::vertex_array vertexArray{};
    opengl::buffer positionBuffer{}; // binding = 0
    opengl::buffer normalBuffer{};   // binding = 1
    opengl::buffer texCoordBuffer{}; // binding = 2
    opengl::buffer indexBuffer{};
};

/*
Auto-generated GLSL attributes and uniforms

Vertex shader:
layout(location = 0) in vec3 inPos;
layout(location = 1) in vec3 inNormal; // Generated if MeshGeometry::normals is not empty
layout(location = 2) in vec2 inTexCoord; // Generated if MeshGeometry::texCoords is not empty

Fragment shader:
layout(location = 0) in vec3 normal;
layout(location = 1) in vec3 fragPos; // MeshPhongMaterial
layout(location = 2) in vec2 texCoord;
layout(location = 0) out vec4 fragColor;

Uniforms:
struct Scene {
    mat4 viewProjMatrix;
    vec3 cameraDir;
    vec3 cameraPos;
};
layout(binding = 0, std140) uniform SceneBlock {
    Scene scene;
};
struct Transform {
    mat4 modelMatrix;
    mat4 normalMatrix;
};
layout(binding = 1, std140) uniform TransformBlock {
    Transform transform;
};
struct DirLight {
    vec3 ambient, diffuse, specular, direction;
};
struct PointLight {
    vec3 ambient;
    float c0;
    vec3 diffuse;
    float c1;
    vec3 specular;
    float c2;
    vec3 position;
};
layout(binding = 2, std430) uniform LightBlock {
    uint numDirLights;
    uint numPointLights;
    layout(offset = 16) PointLight pointLight[4];
    DirLight dirLight[4];
};

MeshBasicMaterial:
struct MeshBasicMaterial {
    vec4 color; // color.w = opacity
};
layout(binding = 3, std430) uniform MaterialBlock {
    MeshBasicMaterial material;
};

MeshPhongMaterial:
struct MeshPhongMaterial {
    vec3 diffuse;
    vec3 specular;
    float shininess;
};
layout(binding = 3, std430) uniform MaterialBlock {
    MeshPhongMaterial material;
};

*/

namespace ShaderStrings {
constexpr auto header = "#version 450 core\n"
                        "#extension GL_ARB_uniform_buffer_object:require\n"
                        "#extension GL_ARB_explicit_attrib_location:require\n"
                        "#extension GL_ARB_explicit_uniform_location:require\n"
                        "#extension GL_NV_uniform_buffer_std430_layout:enable\n";
constexpr auto vertInPos = "layout(location=0)in vec3 inPos;";
constexpr auto vertInNormal = "layout(location=1)in vec3 inNormal;";
constexpr auto vertInTexCoord = "layout(location=2)in vec2 inTexCoord;";
constexpr auto vertOutNormal = "layout(location=0)out vec3 normal;";
constexpr auto vertOutFragPos = "layout(location=1)out vec3 fragPos;";
constexpr auto vertOutTexCoord = "layout(location=2)out vec2 texCoord;";
constexpr auto fragInNormal = "layout(location=0)in vec3 normal;";
constexpr auto fragInFragPos = "layout(location=1)in vec3 fragPos;";
constexpr auto fragInTexCoord = "layout(location=2)in vec2 texCoord;";

} // namespace ShaderStrings

template <> inline void Mesh3D<MeshBasicMaterial>::create(const MeshGeometry &data) {
    using namespace opengl;
    updateMaterial();
    numIndices = data.indices.size();
    positionBuffer.create(sizeof(glm::vec3) * data.positions.size(), data.positions.data());
    indexBuffer.create(4 * numIndices, data.indices.data());
    vertexArray.create({vertex_attribute(0, type_f32, 3, 0, 0)});
    vertexArray.set_vertex_buffer(positionBuffer, 0);
    vertexArray.set_index_buffer(indexBuffer);
}

template <> inline void Mesh3D<MeshPhongMaterial>::create(const MeshGeometry &data) {
    updateMaterial();
    using namespace opengl;
    using namespace ShaderStrings;
    numIndices = data.indices.size();
    positionBuffer.create(sizeof(glm::vec3) * data.positions.size(), data.positions.data());
    normalBuffer.create(sizeof(glm::vec3) * data.normals.size(), data.normals.data());
    indexBuffer.create(4 * numIndices, data.indices.data());
    // inPos: location = 0, binding = 0 (positionBuffer)
    // inNormal: location = 1, binding = 1 (normalBuffer)
    vertexArray.create({
        vertex_attribute(0, type_f32, 3, 0, 0),
        vertex_attribute(1, type_f32, 3, 0, 1),
    });
    vertexArray.set_vertex_buffer(positionBuffer, sizeof(glm::vec3), 0, 0);
    vertexArray.set_vertex_buffer(normalBuffer, sizeof(glm::vec3), 0, 1);
    vertexArray.set_index_buffer(indexBuffer);

    std::string vsSource(header), fsSource(header);
    vsSource += R"(
layout(location = 0) in vec3 inPos;
layout(location = 1) in vec3 inNormal;
layout(location = 0) out vec3 normal;
layout(location = 1) out vec3 fragPos;
struct Scene {
    mat4 viewProjectionMatrix;
    vec3 cameraDir;
    vec3 cameraPos;
};
layout(binding = 0, std140) uniform SceneBlock {
    Scene scene;
};
struct Transform {
    mat4 modelMatrix;
    mat4 normalMatrix;
};
layout(binding = 1, std140) uniform TransformBlock {
    Transform transform;
};
void main() {
    const mat4 worldMatrix = scene.viewProjectionMatrix * transform.modelMatrix;
    gl_Position = worldMatrix * vec4(inPos, 1.0f);
    normal = vec3(transform.normalMatrix * vec4(inNormal, 1.0f));
    fragPos = vec3(transform.modelMatrix * vec4(inPos, 1.0f));
})";
    fsSource += R"(
layout(location = 0) in vec3 normal;
layout(location = 1) in vec3 fragPos;
layout(location = 0) out vec4 fragColor;
struct Scene {
    mat4 viewProjectionMatrix;
    vec3 cameraDir;
    vec3 cameraPos;
};
layout(binding = 0, std140) uniform SceneBlock {
    Scene scene;
};
struct DirLight {
    vec3 ambient, diffuse, specular, direction;
};
struct PointLight {
    vec3 ambient;
    float c0;
    vec3 diffuse;
    float c1;
    vec3 specular;
    float c2;
    vec3 position;
};
layout(binding = 2, std430) uniform LightBlock {
    uint numDirLights;
    uint numPointLights;
    layout(offset = 16) DirLight dirLight[4];
    PointLight pointLight[4];
};
struct MeshPhongMaterial {
    vec3 diffuse;
    vec3 specular;
    float shininess;
};
layout(binding = 3, std430) uniform MaterialBlock {
    MeshPhongMaterial material;
};

void main() {
    vec3 result = vec3(0.0);
    const vec3 viewDir = normalize(scene.cameraPos - fragPos);
    for (uint i = 0; i < numDirLights; ++i) {
        vec3 ambient = dirLight[i].ambient;
        vec3 lightDir = -normalize(dirLight[i].direction);
        vec3 diffuse = max(dot(lightDir, normal), 0.0f) * dirLight[i].diffuse * material.diffuse;
        vec3 reflectDir = reflect(-lightDir, normal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0f), material.shininess);
        vec3 specular = spec * dirLight[i].specular * material.specular;
        result += (ambient + diffuse + specular);
    }
    for (uint i = 0; i < numPointLights; ++i) {
        float dis = distance(pointLight[i].position, fragPos);
        float attenuation = pointLight[i].c0 + pointLight[i].c1 * dis + pointLight[i].c2 * dis * dis;
        vec3 ambient = pointLight[i].ambient;
        vec3 lightDir = normalize(pointLight[i].position - fragPos);
        vec3 diffuse = max(dot(lightDir, normal), 0.0f) * pointLight[i].diffuse;
        vec3 reflectDir = reflect(-lightDir, normal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0f), material.shininess);
        vec3 specular = spec * pointLight[i].specular * material.specular;
        result += (ambient + diffuse + specular) / attenuation;
    }
    fragColor = vec4(result, 1.0f);
})";
    program.create(
        {shader(shader::vertex_shader, vsSource.c_str()), shader(shader::fragment_shader, fsSource.c_str())});
    pipeline.depthTest.enabled = true;
}

template <> void Mesh3D<MeshPhongMaterial>::update() { Entity3D::update(); }
template <> void Mesh3D<MeshPhongMaterial>::render() {
    program.use();
    vertexArray.bind();
    pipeline.activate();
    transformBuffer.bind_to(opengl::buffer::uniform_buffer, 1);
    materialBuffer.bind_to(opengl::buffer::uniform_buffer, 3);
    pipeline.drawIndexed(numIndices);
    pipeline.deactivate();
}
